package edu.hawaii.ics.pjdb.result_transport;

import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.apache.thrift.TException;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.server.THsHaServer;
import org.apache.thrift.server.TServer;
import org.apache.thrift.server.TSimpleServer;
import org.apache.thrift.transport.TNonblockingServerSocket;
import org.apache.thrift.transport.TServerSocket;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import au.com.bytecode.opencsv.CSVReader;

import edu.hawaii.ics.pjdb.types.RemoteAddress;
import edu.hawaii.ics.pjdb.types.Tuple;
import edu.hawaii.ics.pjdb.types.TupleService;
import edu.hawaii.ics.pjdb.util.ThriftHelper;

public class TupleServiceHelper {
    private static final Logger logger = LoggerFactory
            .getLogger(TupleServiceHelper.class);

    /**
     * Number of tuples to buffer before sending them off.
     */
    private static final int tupleBufferSize = 100;

    /**
     * Number of milliseconds to wait before giving up when connecting to a
     * server
     */
    private static final int connectionTimeout = 30000;

    /**
     * Number of milliseconds to wait before retrying a failed connection
     */
    private static final int retryPeriod = 1000;

    /**
     * Serializes the row as thrift Tuple. This is probably not the best way to
     * do this...
     * 
     * @param results
     * @return
     * @throws SQLException
     */
    public static Tuple serializeResult(ResultSet results) throws SQLException {
        int numCols = results.getMetaData().getColumnCount();
        List<String> entries = new ArrayList<String>(numCols);

        for (int i = 1; i <= numCols; i++) {
            entries.add(results.getString(i));
        }

        return new Tuple(entries);
    }

    /**
     * Concatenates two tuples together
     * 
     * @param t1
     * @param t2
     * @return
     */
    public static Tuple join(Tuple t1, Tuple t2) {
        List<String> entries = new ArrayList<String>(t1.getEntriesSize()
                + t2.getEntriesSize());
        entries.addAll(t1.getEntries());
        entries.addAll(t2.getEntries());

        return new Tuple(entries);
    }

    /**
     * Same as {@link #submitTuples(TupleService.Client, ResultSet)}, but calls
     * {@link #join(Tuple, Tuple)} on every tuple that is returned.
     * 
     * @param client
     * @param results
     * @param joinTuple
     * @throws SQLException
     * @throws TException
     */
    public static void submitTuples(TupleService.Client client,
            ResultSet results, Tuple joinTuple) throws SQLException, TException {
        List<Tuple> buffer = new ArrayList<Tuple>(tupleBufferSize);

        while (results.next()) {
            buffer.add(join(serializeResult(results), joinTuple));

            if (buffer.size() >= tupleBufferSize) {
                client.submitTuples(buffer);
                buffer.clear();
            }
        }

        if (buffer.size() > 0) {
            client.submitTuples(buffer);
        }
    }

    /**
     * Submits all tuples in the specified {@link ResultSet} through the given
     * thrift client. Does not call {@link TupleService.Client#finished()} when
     * done.
     * 
     * @param client
     * @param results
     * @throws SQLException
     * @throws TException
     */
    public static void submitTuples(TupleService.Client client,
            ResultSet results) throws SQLException, TException {
        List<Tuple> buffer = new ArrayList<Tuple>(tupleBufferSize);

        while (results.next()) {
            buffer.add(serializeResult(results));

            if (buffer.size() >= tupleBufferSize) {
                client.submitTuples(buffer);
                buffer.clear();
            }
        }

        if (buffer.size() > 0) {
            client.submitTuples(buffer);
        }
    }

    /**
     * Sends the tuples contained in results to the tuple service running on the
     * specified machine/port.
     * 
     * @param addr
     * @param results
     * @throws TTransportException
     * @throws SQLException
     */
    public static void submitTuples(RemoteAddress addr, ResultSet results)
            throws TException, SQLException {
        TupleService.Client client = createClientWithRetry(addr);

        submitTuples(client, results);

        try {
            client.finished();
            client.getOutputProtocol().getTransport().close();
        } catch (TTransportException e) {
            // We expect this to happen.
        }
    }

    /**
     * Reads tuples from the given CSV file and sends the tuple using the
     * client.
     * 
     * @param client
     * @param csv
     * @param header
     * @throws SQLException
     * @throws TException
     * @throws IOException
     * @throws FileNotFoundException
     */
    public static void submitTuples(TupleService.Client client, String csv,
            boolean header) throws SQLException, TException, IOException,
            FileNotFoundException {
        List<Tuple> buffer = new ArrayList<Tuple>(tupleBufferSize);
        CSVReader reader = new CSVReader(new FileReader(csv));
        String[] nextLine;

        // Skip over header.
        if (header) {
            reader.readNext();
        }

        nextLine = reader.readNext();
        while (nextLine != null) {
            buffer.add(new Tuple(Arrays.asList(nextLine)));
            if (buffer.size() >= tupleBufferSize) {
                client.submitTuples(buffer);
                buffer.clear();
            }
            nextLine = reader.readNext();
        }

        if (buffer.size() > 0) {
            client.submitTuples(buffer);
        }
    }

    /**
     * Reads tuples from the given CSV file and sends them to the worker with
     * the given RemoteAddress.
     * 
     * @param addr
     * @param csv
     * @param header
     * @throws TException
     * @throws SQLException
     * @throws FileNotFoundException
     * @throws IOException
     */
    public static void submitTuples(RemoteAddress addr, String csv,
            boolean header) throws TException, SQLException,
            FileNotFoundException, IOException {
        TupleService.Client client = createClientWithRetry(addr);

        submitTuples(client, csv, header);

        try {
            client.finished();
            client.getOutputProtocol().getTransport().close();
        } catch (TTransportException e) {
            // We expect this to happen.
        }
    }

    /**
     * Sends the list of Tuples to the worker.
     * 
     * @param client
     * @param csv
     * @param header
     * @throws SQLException
     * @throws TException
     * @throws IOException
     * @throws FileNotFoundException
     */
    public static void submitTuples(TupleService.Client client,
            List<Tuple> tuples) throws TException {
        List<Tuple> buffer = new ArrayList<Tuple>(tupleBufferSize);

        // Submit tuples in buffer sized chunks
        while (tuples.size() > tupleBufferSize) {
            buffer = tuples.subList(0, tupleBufferSize - 1);
            tuples.removeAll(buffer);
            client.submitTuples(buffer);
        }

        // Tuples smaller than max buffer size. Submit entire tuple list.
        client.submitTuples(tuples);
    }

    /**
     * Sends the list of Tuples to the worker with the given RemoteAddress.
     * 
     * @param addr
     * @param csv
     * @param header
     * @throws TException
     * @throws SQLException
     * @throws FileNotFoundException
     * @throws IOException
     */
    public static void submitTuples(RemoteAddress addr, List<Tuple> tuples)
            throws TException {
        TupleService.Client client = createClientWithRetry(addr);

        submitTuples(client, tuples);

        try {
            client.finished();
            client.getOutputProtocol().getTransport().close();
        } catch (TTransportException e) {
            // We expect this to happen.
        }
    }

    /**
     * Creates a client connecting to the remote server specified by addr.
     * 
     * @param addr
     * @return
     * @throws TTransportException
     */
    public static TupleService.Client createClient(RemoteAddress addr)
            throws TTransportException {
        TSocket socket = new TSocket(addr.getAddress(), addr.getPort());
        TTransport transport = ThriftHelper.getTransportFactory().getTransport(
                socket);
        TProtocol protocol = ThriftHelper.getProtocolFactory().getProtocol(
                transport);

        // Open the transport
        transport.open();

        return new TupleService.Client(protocol);
    }

    /**
     * Tries very diligently to create a connection to a tuple server. It will
     * try again every second for the amount of time specified in
     * {@link #connectionTimeout}. If it still fails after that period of time,
     * it throws the same exception that {@link #createClient(RemoteAddress)}
     * would.
     * 
     * @param addr
     * @return
     * @throws TTransportException
     * @throws InterruptedException
     */
    public static TupleService.Client createClientWithRetry(RemoteAddress addr)
            throws TTransportException {
        long startTime = System.currentTimeMillis();

        // Keep track of the exception that we should throw.
        TTransportException theException = null;

        while ((System.currentTimeMillis() - startTime) < connectionTimeout) {
            try {
                return createClient(addr);
            } catch (TTransportException e) {
                // Wait for a second and then retry
                try {
                    Thread.sleep(retryPeriod);
                } catch (InterruptedException ie) {
                    throw e;
                }

                theException = e;
                logger.info(
                        "Failed to connect after {} milliseconds... retrying in {}ms.",
                        System.currentTimeMillis() - startTime, retryPeriod);
            }
        }

        throw theException;
    }

    /**
     * Create a tuple server that listens for results on the specified port. It
     * also takes a {@link TupleListener} which should be responsible for
     * handling the results.
     * 
     * NOTE: in order to start serving, you must call {@link TServer#serve()} on
     * the value that this method returns to begin listening for results.
     * 
     * The server that is returned should be non-blocking.
     * 
     * @param port
     * @param listener
     * @return
     * @throws TTransportException
     */
    public static TServer createServer(int port, TupleListener listener)
            throws TTransportException {
        TupleReceiver handler = new TupleReceiver(listener);
        TupleService.Processor<TupleService.Iface> processor = new TupleService.Processor<TupleService.Iface>(
                handler);
        TServer server = new THsHaServer(new THsHaServer.Args(
                new TNonblockingServerSocket(port)).processor(processor)
                .transportFactory(ThriftHelper.getTransportFactory())
                .protocolFactory(ThriftHelper.getProtocolFactory()));
        handler.setServer(server);

        return server;
    }

    /**
     * Create a simple, non-multithreaded tuple server that listens for results
     * on the specified port. It also takes a {@link TupleListener} which should
     * be responsible for handling the results.
     * 
     * NOTE: in order to start serving, you must call {@link TServer#serve()} on
     * the value that this method returns to begin listening for results.
     * 
     * @param port
     * @param listener
     * @return
     * @throws TTransportException
     */
    public static TServer createNonThreadedServer(int port,
            TupleListener listener) throws TTransportException {
        TupleReceiver handler = new TupleReceiver(listener);
        TupleService.Processor<TupleService.Iface> processor = new TupleService.Processor<TupleService.Iface>(
                handler);
        TServer server = new TSimpleServer(new TServer.Args(new TServerSocket(
                port)).processor(processor)
                .transportFactory(ThriftHelper.getTransportFactory())
                .protocolFactory(ThriftHelper.getProtocolFactory()));
        handler.setServer(server);

        return server;
    }
}
